import matplotlib.pyplot as plt
from collections import defaultdict, Counter
import argparse

# 从文件中读取并解析区间数据的函数
def read_kmer_data(file_path):
    interval_data = defaultdict(dict)
    current_interval = None
    
    with open(file_path, 'r') as file:
        for line in file:
            line = line.strip()
            if line.startswith("Interval:"):
                current_interval = line.split("Interval: ")[1].strip()
            elif line and current_interval:
                parts = line.split(":")
                if len(parts) == 2:
                    kmer = parts[0].strip()
                    try:
                        count = int(parts[1].strip())
                        # 仅处理kmer长度等于5的数据
                        if len(kmer) == 5:
                            interval_data[current_interval][kmer] = count
                    except ValueError:
                        # 跳过无法转换为整数的行
                        pass

    return interval_data

# 根据用户输入选择区间并绘制图表的函数
def plot_top_kmers(interval_data, selected_interval, top_n=10):
    if selected_interval not in interval_data:
        print(f"区间 '{selected_interval}' 不存在。请检查输入。")
        return
    
    # 获取选择区间的kmer频率数据
    frequency_data = interval_data[selected_interval]
    
    # 使用Counter获取出现次数最多的top_n个kmer
    top_kmers = Counter(frequency_data).most_common(top_n)
    if not top_kmers:
        print(f"区间 '{selected_interval}' 中没有长度为5的kmer。")
        return

    top_kmers_labels, top_kmers_counts = zip(*top_kmers)

    # 绘制柱状图
    plt.figure(figsize=(10, 6))
    plt.bar(top_kmers_labels, top_kmers_counts, color='skyblue')
    plt.xlabel('K-mer')
    plt.ylabel('Frequency')
    plt.title(f'Top {top_n} K-mers in Interval {selected_interval}')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

# 主函数，用于整体控制逻辑
def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="选择性读取kmer数据并绘制图表")
    parser.add_argument("file_path", help="输入的kmer统计文件路径")
    args = parser.parse_args()

    # 读取kmer数据
    interval_data = read_kmer_data(args.file_path)
    
    # 打印可用的区间名称
    print("可用的区间名称:")
    for interval in interval_data.keys():
        print(interval)
    
    # 用户选择区间
    selected_interval = input("请输入要分析的区间名称：")
    
    # 绘制图表
    plot_top_kmers(interval_data, selected_interval)

if __name__ == "__main__":
    main()
